"""MongoDB客户端 - 长期情节记忆存储"""

import asyncio
from typing import Dict, Any, Optional, List, Union
from datetime import datetime, timedelta
from bson import ObjectId
from loguru import logger

try:
    from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase, AsyncIOMotorCollection
except ImportError:
    raise ImportError("请安装motor: pip install motor")

from src.config import get_database_url


class MongoDBClient:
    """MongoDB客户端 - 负责长期情节记忆存储"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.mongodb_config = config.get('databases', {}).get('mongodb', {})
        self.client: Optional[AsyncIOMotorClient] = None
        self.database: Optional[AsyncIOMotorDatabase] = None
        self._connected = False
    
    async def connect(self):
        """连接MongoDB"""
        if self._connected:
            return
        
        try:
            mongodb_url = get_database_url(self.config, 'mongodb')
            
            self.client = AsyncIOMotorClient(
                mongodb_url,
                serverSelectionTimeoutMS=self.mongodb_config.get('timeout', 5000),
                connectTimeoutMS=self.mongodb_config.get('connect_timeout', 5000),
                maxPoolSize=self.mongodb_config.get('max_pool_size', 10),
                minPoolSize=self.mongodb_config.get('min_pool_size', 1)
            )
            
            # 获取数据库
            db_name = self.mongodb_config.get('database', 'memory_ai')
            self.database = self.client[db_name]
            
            # 测试连接
            await self.client.admin.command('ping')
            
            # 创建索引
            await self._create_indexes()
            
            self._connected = True
            logger.info(f"MongoDB连接成功: {db_name}")
            
        except Exception as e:
            logger.error(f"MongoDB连接失败: {e}")
            raise
    
    async def close(self):
        """关闭MongoDB连接"""
        if self.client:
            self.client.close()
            self._connected = False
            logger.info("MongoDB连接已关闭")
    
    async def ping(self) -> bool:
        """检查连接状态"""
        try:
            if self.client:
                await self.client.admin.command('ping')
                return True
            return False
        except Exception:
            return False
    
    async def _create_indexes(self):
        """创建必要的索引"""
        try:
            # 对话集合索引
            conversations = self.database.conversations
            await conversations.create_index("session_id")
            await conversations.create_index("user_id")
            await conversations.create_index("timestamp")
            await conversations.create_index([("session_id", 1), ("timestamp", -1)])
            # 新增：按业务主键与常用查询字段建立索引
            await conversations.create_index("conversation_id")
            await conversations.create_index("updated_at")
            await conversations.create_index("status")
            await conversations.create_index([("user_id", 1), ("updated_at", -1)])

            # 记忆集合索引
            memories = self.database.memories
            await memories.create_index("user_id")
            await memories.create_index("memory_type")
            await memories.create_index("timestamp")
            await memories.create_index("tags")
            await memories.create_index([("user_id", 1), ("memory_type", 1)])

            # 事件集合索引
            events = self.database.events
            await events.create_index("user_id")
            await events.create_index("event_type")
            await events.create_index("timestamp")
            await events.create_index([("user_id", 1), ("timestamp", -1)])

            # 知识图谱节点集合索引
            knowledge_nodes = self.database.knowledge_nodes
            await knowledge_nodes.create_index("node_id")
            await knowledge_nodes.create_index("node_type")
            await knowledge_nodes.create_index("user_id")

            # 新增：对话消息集合索引
            conversation_messages = self.database.conversation_messages
            await conversation_messages.create_index("conversation_id")
            await conversation_messages.create_index("timestamp")
            await conversation_messages.create_index("role")
            await conversation_messages.create_index([("conversation_id", 1), ("timestamp", 1)])

            # 文档集合索引
            documents = self.database.documents
            await documents.create_index("document_id")
            await documents.create_index("user_id")
            await documents.create_index("status")
            await documents.create_index("updated_at")
            await documents.create_index("created_at")
            await documents.create_index([("user_id", 1), ("updated_at", -1)])
            await documents.create_index([("document_id", 1), ("status", 1)])

            logger.info("MongoDB索引创建完成")

        except Exception as e:
            logger.error(f"创建MongoDB索引失败: {e}")
    
    def get_collection(self, name: str) -> AsyncIOMotorCollection:
        """获取集合"""
        if self.database is None:
            raise RuntimeError("MongoDB未连接")
        return self.database[name]
    
    # 对话记录操作
    async def store_conversation(self, conversation_data: Dict[str, Any]) -> str:
        """存储对话记录"""
        try:
            collection = self.get_collection('conversations')
            
            # 添加时间戳
            conversation_data['timestamp'] = datetime.utcnow()
            conversation_data['created_at'] = datetime.utcnow()
            conversation_data['updated_at'] = datetime.utcnow()
            
            result = await collection.insert_one(conversation_data)
            logger.debug(f"对话记录已存储: {result.inserted_id}")
            return str(result.inserted_id)
            
        except Exception as e:
            logger.error(f"存储对话记录失败: {e}")
            raise
    
    async def get_conversation(self, conversation_id: str) -> Optional[Dict[str, Any]]:
        """获取对话记录（优先按业务ID conversation_id，其次回退到 _id）"""
        try:
            collection = self.get_collection('conversations')
            # 先按业务ID查询
            result = await collection.find_one({"conversation_id": conversation_id})

            # 若未命中，尝试按 _id 查询
            if not result:
                try:
                    result = await collection.find_one({"_id": ObjectId(conversation_id)})
                except Exception:
                    result = None

            if result and "_id" in result:
                result["_id"] = str(result["_id"])

            return result

        except Exception as e:
            logger.error(f"获取对话记录失败: {e}")
            return None
    
    async def get_conversations_by_session(self, session_id: str, limit: int = 50) -> List[Dict[str, Any]]:
        """根据会话ID获取对话记录"""
        try:
            collection = self.get_collection('conversations')
            cursor = collection.find(
                {"session_id": session_id}
            ).sort("timestamp", -1).limit(limit)
            
            conversations = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                conversations.append(doc)
            
            return conversations
            
        except Exception as e:
            logger.error(f"获取会话对话记录失败: {e}")
            return []
    
    async def get_conversations_by_user(self, user_id: str, limit: int = 100, skip: int = 0) -> List[Dict[str, Any]]:
        """根据用户ID获取对话记录（按timestamp倒序）"""
        try:
            collection = self.get_collection('conversations')
            cursor = collection.find(
                {"user_id": user_id}
            ).sort("timestamp", -1).skip(skip).limit(limit)
            
            conversations = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                conversations.append(doc)
            
            return conversations
            
        except Exception as e:
            logger.error(f"获取用户对话记录失败: {e}")
            return []
    
    # 记忆操作
    async def store_memory(self, memory_data: Dict[str, Any]) -> str:
        """存储记忆"""
        try:
            collection = self.get_collection('memories')
            
            # 添加时间戳
            memory_data['timestamp'] = datetime.utcnow()
            memory_data['created_at'] = datetime.utcnow()
            memory_data['updated_at'] = datetime.utcnow()
            
            result = await collection.insert_one(memory_data)
            logger.debug(f"记忆已存储: {result.inserted_id}")
            return str(result.inserted_id)
            
        except Exception as e:
            logger.error(f"存储记忆失败: {e}")
            raise
    
    async def get_memories_by_user(self, user_id: str, memory_type: Optional[str] = None, limit: int = 50) -> List[Dict[str, Any]]:
        """根据用户ID获取记忆"""
        try:
            collection = self.get_collection('memories')
            
            query = {"user_id": user_id}
            if memory_type:
                query["memory_type"] = memory_type
            
            cursor = collection.find(query).sort("timestamp", -1).limit(limit)
            
            memories = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                memories.append(doc)
            
            return memories
            
        except Exception as e:
            logger.error(f"获取用户记忆失败: {e}")
            return []
    
    async def search_memories(self, user_id: str, query: str, memory_type: Optional[str] = None, limit: int = 20) -> List[Dict[str, Any]]:
        """搜索记忆"""
        try:
            collection = self.get_collection('memories')
            
            # 构建搜索条件
            search_query = {
                "user_id": user_id,
                "$or": [
                    {"content": {"$regex": query, "$options": "i"}},
                    {"title": {"$regex": query, "$options": "i"}},
                    {"tags": {"$regex": query, "$options": "i"}}
                ]
            }
            
            if memory_type:
                search_query["memory_type"] = memory_type
            
            cursor = collection.find(search_query).sort("timestamp", -1).limit(limit)
            
            memories = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                memories.append(doc)
            
            return memories
            
        except Exception as e:
            logger.error(f"搜索记忆失败: {e}")
            return []
    
    # 事件操作
    async def store_event(self, event_data: Dict[str, Any]) -> str:
        """存储事件"""
        try:
            collection = self.get_collection('events')
            
            # 添加时间戳
            event_data['timestamp'] = datetime.utcnow()
            event_data['created_at'] = datetime.utcnow()
            
            result = await collection.insert_one(event_data)
            logger.debug(f"事件已存储: {result.inserted_id}")
            return str(result.inserted_id)
            
        except Exception as e:
            logger.error(f"存储事件失败: {e}")
            raise
    
    async def get_events_by_user(self, user_id: str, event_type: Optional[str] = None, 
                                start_time: Optional[datetime] = None, 
                                end_time: Optional[datetime] = None, 
                                limit: int = 100) -> List[Dict[str, Any]]:
        """根据用户ID获取事件"""
        try:
            collection = self.get_collection('events')
            
            query = {"user_id": user_id}
            
            if event_type:
                query["event_type"] = event_type
            
            if start_time or end_time:
                time_query = {}
                if start_time:
                    time_query["$gte"] = start_time
                if end_time:
                    time_query["$lte"] = end_time
                query["timestamp"] = time_query
            
            cursor = collection.find(query).sort("timestamp", -1).limit(limit)
            
            events = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                events.append(doc)
            
            return events
            
        except Exception as e:
            logger.error(f"获取用户事件失败: {e}")
            return []
    
    # 知识图谱节点操作
    async def store_knowledge_node(self, node_data: Dict[str, Any]) -> str:
        """存储知识图谱节点"""
        try:
            collection = self.get_collection('knowledge_nodes')
            
            # 添加时间戳
            node_data['timestamp'] = datetime.utcnow()
            node_data['created_at'] = datetime.utcnow()
            node_data['updated_at'] = datetime.utcnow()
            
            result = await collection.insert_one(node_data)
            logger.debug(f"知识节点已存储: {result.inserted_id}")
            return str(result.inserted_id)
            
        except Exception as e:
            logger.error(f"存储知识节点失败: {e}")
            raise
    
    async def get_knowledge_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        """获取知识图谱节点"""
        try:
            collection = self.get_collection('knowledge_nodes')
            result = await collection.find_one({"node_id": node_id})
            
            if result:
                result['_id'] = str(result['_id'])
            
            return result
            
        except Exception as e:
            logger.error(f"获取知识节点失败: {e}")
            return None
    
    # 统计操作
    async def get_user_statistics(self, user_id: str) -> Dict[str, Any]:
        """获取用户统计信息"""
        try:
            stats = {}
            
            # 对话统计
            conversations_count = await self.get_collection('conversations').count_documents({"user_id": user_id})
            stats['conversations_count'] = conversations_count
            
            # 记忆统计
            memories_count = await self.get_collection('memories').count_documents({"user_id": user_id})
            stats['memories_count'] = memories_count
            
            # 事件统计
            events_count = await self.get_collection('events').count_documents({"user_id": user_id})
            stats['events_count'] = events_count
            
            # 知识节点统计
            nodes_count = await self.get_collection('knowledge_nodes').count_documents({"user_id": user_id})
            stats['knowledge_nodes_count'] = nodes_count
            
            return stats
            
        except Exception as e:
            logger.error(f"获取用户统计失败: {e}")
            return {}
    
    async def get_user_stats(self, user_id: str) -> Dict[str, Any]:
        """获取用户统计信息（别名方法）"""
        return await self.get_user_statistics(user_id)
    
    async def update_user_stats(self, user_id: str, updates: Dict[str, Any]) -> bool:
        """更新用户统计信息"""
        try:
            collection = self.get_collection('user_stats')
            
            # 添加更新时间
            updates['updated_at'] = datetime.utcnow()
            
            # 使用upsert操作，如果不存在则创建
            result = await collection.update_one(
                {"user_id": user_id},
                {"$set": updates},
                upsert=True
            )
            
            return result.modified_count > 0 or result.upserted_id is not None
            
        except Exception as e:
            logger.error(f"更新用户统计失败: {e}")
            return False
    
    # 数据清理
    async def cleanup_old_data(self, days: int = 90):
        """清理旧数据"""
        try:
            cutoff_date = datetime.utcnow() - timedelta(days=days)
            
            # 清理旧对话
            conversations_result = await self.get_collection('conversations').delete_many(
                {"timestamp": {"$lt": cutoff_date}}
            )
            
            # 清理旧事件
            events_result = await self.get_collection('events').delete_many(
                {"timestamp": {"$lt": cutoff_date}}
            )
            
            logger.info(f"清理了 {conversations_result.deleted_count} 条对话记录和 {events_result.deleted_count} 条事件记录")
            
        except Exception as e:
            logger.error(f"清理旧数据失败: {e}")
    
    async def store_conversation_message(self, conversation_id: str, message_data: Dict[str, Any]) -> str:
        """存储对话消息到 conversation_messages 集合"""
        try:
            collection = self.get_collection('conversation_messages')
            # 确保会话ID存在
            message_data['conversation_id'] = conversation_id
            # 确保时间戳
            if 'timestamp' not in message_data or message_data['timestamp'] is None:
                message_data['timestamp'] = datetime.utcnow()
            # 元数据兜底
            if 'metadata' not in message_data or message_data['metadata'] is None:
                message_data['metadata'] = {}

            result = await collection.insert_one(message_data)
            logger.debug(f"对话消息已存储: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"存储对话消息失败: {e}")
            raise
    
    async def get_conversation_messages(self, conversation_id: str, limit: Optional[int] = None, offset: int = 0) -> List[Dict[str, Any]]:
        """获取对话消息，按时间正序返回以便还原历史"""
        try:
            collection = self.get_collection('conversation_messages')
            cursor = collection.find({
                'conversation_id': conversation_id
            }).sort('timestamp', 1)

            if offset:
                cursor = cursor.skip(offset)
            if limit is not None:
                cursor = cursor.limit(limit)

            messages: List[Dict[str, Any]] = []
            async for doc in cursor:
                if '_id' in doc:
                    doc['_id'] = str(doc['_id'])
                # 标准化字段
                doc.setdefault('metadata', {})
                messages.append(doc)
            return messages
        except Exception as e:
            logger.error(f"获取对话消息失败: {e}")
            return []
    
    async def update_conversation(self, conversation_id: str, updates: Dict[str, Any]) -> bool:
        """更新对话信息（按 conversation_id 更新，必要时兼容 _id）"""
        try:
            collection = self.get_collection('conversations')
            # 自动更新时间
            updates = dict(updates or {})
            updates['updated_at'] = updates.get('updated_at', datetime.utcnow())

            # 先尝试按业务ID更新
            result = await collection.update_one(
                {"conversation_id": conversation_id},
                {"$set": updates}
            )

            if result.matched_count == 0:
                # 回退按 _id
                try:
                    result = await collection.update_one(
                        {"_id": ObjectId(conversation_id)},
                        {"$set": updates}
                    )
                except Exception:
                    pass

            return result.modified_count > 0 or result.matched_count > 0
        except Exception as e:
            logger.error(f"更新对话失败: {e}")
            return False
    
    async def get_user_conversations(self, user_id: str, limit: int = 20, status: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取用户的对话列表，支持按状态过滤，按 updated_at 倒序"""
        try:
            collection = self.get_collection('conversations')
            query: Dict[str, Any] = {"user_id": user_id}
            if status:
                query["status"] = status
            cursor = collection.find(query).sort("updated_at", -1).limit(limit)

            items: List[Dict[str, Any]] = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                items.append(doc)
            return items
        except Exception as e:
            logger.error(f"获取用户对话失败: {e}")
            return []
    
    async def get_conversations_before_date(self, cutoff_date: datetime) -> List[Dict[str, Any]]:
        """获取在指定日期之前更新的对话（用于清理）"""
        try:
            collection = self.get_collection('conversations')
            cursor = collection.find({
                "$or": [
                    {"updated_at": {"$lt": cutoff_date}},
                    {"timestamp": {"$lt": cutoff_date}}
                ]
            }).sort("updated_at", 1)

            items: List[Dict[str, Any]] = []
            async for doc in cursor:
                doc['_id'] = str(doc['_id'])
                items.append(doc)
            return items
        except Exception as e:
            logger.error(f"获取旧对话失败: {e}")
            return []
    
    async def get_conversation_stats(self) -> Dict[str, Any]:
        """获取对话统计信息"""
        try:
            conversations = self.get_collection('conversations')
            total = await conversations.count_documents({})
            active = await conversations.count_documents({"status": "active"})
            closed = await conversations.count_documents({"status": "closed"})

            # 近7日新增
            seven_days_ago = datetime.utcnow() - timedelta(days=7)
            recent = await conversations.count_documents({"created_at": {"$gte": seven_days_ago}})

            # 消息数量统计（可能较慢，生产可改为增量维护）
            messages_col = self.get_collection('conversation_messages')
            messages_total = await messages_col.count_documents({})

            return {
                "total_conversations": total,
                "active_conversations": active,
                "closed_conversations": closed,
                "recent_conversations_7d": recent,
                "messages_total": messages_total,
            }
        except Exception as e:
            logger.error(f"获取对话统计失败: {e}")
            return {}
    
    async def delete_conversation(self, conversation_id: str) -> bool:
        """删除对话及其消息"""
        try:
            conversations = self.get_collection('conversations')
            messages = self.get_collection('conversation_messages')

            # 先删除消息
            await messages.delete_many({"conversation_id": conversation_id})

            # 删除对话（优先按业务ID）
            result = await conversations.delete_one({"conversation_id": conversation_id})
            if result.deleted_count == 0:
                # 回退使用 _id
                try:
                    result = await conversations.delete_one({"_id": ObjectId(conversation_id)})
                except Exception:
                    pass

            return result.deleted_count > 0
        except Exception as e:
            logger.error(f"删除对话失败: {e}")
            return False

    # ==================== 文档相关方法 ====================
    
    async def store_document(self, document_data: Dict[str, Any]) -> str:
        """存储文档"""
        try:
            documents = self.get_collection('documents')
            
            # 确保必要字段
            if 'created_at' not in document_data:
                document_data['created_at'] = datetime.utcnow()
            if 'updated_at' not in document_data:
                document_data['updated_at'] = datetime.utcnow()
            if 'status' not in document_data:
                document_data['status'] = 'pending'
            
            result = await documents.insert_one(document_data)
            logger.debug(f"文档存储成功: {result.inserted_id}")
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"存储文档失败: {e}")
            raise
    
    async def get_document(self, document_id: str) -> Optional[Dict[str, Any]]:
        """获取文档"""
        try:
            documents = self.get_collection('documents')
            
            # 优先按业务ID查询
            document = await documents.find_one({"document_id": document_id})
            if not document:
                # 回退使用 _id
                try:
                    document = await documents.find_one({"_id": ObjectId(document_id)})
                except Exception:
                    pass
            
            if document:
                document['_id'] = str(document['_id'])
            return document
        except Exception as e:
            logger.error(f"获取文档失败: {e}")
            return None
    
    async def update_document(self, document_id: str, updates: Dict[str, Any]) -> bool:
        """更新文档"""
        try:
            documents = self.get_collection('documents')
            
            # 添加更新时间
            updates['updated_at'] = datetime.utcnow()
            
            # 优先按业务ID更新
            result = await documents.update_one(
                {"document_id": document_id},
                {"$set": updates}
            )
            
            if result.matched_count == 0:
                # 回退使用 _id
                try:
                    result = await documents.update_one(
                        {"_id": ObjectId(document_id)},
                        {"$set": updates}
                    )
                except Exception:
                    pass
            
            return result.modified_count > 0
        except Exception as e:
            logger.error(f"更新文档失败: {e}")
            return False
    
    async def delete_document(self, document_id: str) -> bool:
        """删除文档"""
        try:
            documents = self.get_collection('documents')
            
            # 优先按业务ID删除
            result = await documents.delete_one({"document_id": document_id})
            if result.deleted_count == 0:
                # 回退使用 _id
                try:
                    result = await documents.delete_one({"_id": ObjectId(document_id)})
                except Exception:
                    pass
            
            return result.deleted_count > 0
        except Exception as e:
            logger.error(f"删除文档失败: {e}")
            return False
    
    async def get_user_documents(self, user_id: str, limit: int = 50, skip: int = 0, status: Optional[str] = None) -> List[Dict[str, Any]]:
        """获取用户文档列表"""
        try:
            documents = self.get_collection('documents')
            
            query = {"user_id": user_id}
            if status:
                query["status"] = status
            
            cursor = documents.find(query).sort("updated_at", -1).skip(skip).limit(limit)
            docs = await cursor.to_list(length=limit)
            
            # 转换 ObjectId
            for doc in docs:
                doc['_id'] = str(doc['_id'])
            
            return docs
        except Exception as e:
            logger.error(f"获取用户文档失败: {e}")
            return []
    
    async def get_document_stats(self, user_id: Optional[str] = None) -> Dict[str, Any]:
        """获取文档统计信息"""
        try:
            documents = self.get_collection('documents')
            
            query = {}
            if user_id:
                query["user_id"] = user_id
            
            total = await documents.count_documents(query)
            
            # 按状态统计
            pipeline = [
                {"$match": query},
                {"$group": {
                    "_id": "$status",
                    "count": {"$sum": 1}
                }}
            ]
            
            status_stats = {}
            async for result in documents.aggregate(pipeline):
                status_stats[result["_id"]] = result["count"]
            
            # 最近7天的文档数量
            recent_query = query.copy()
            recent_query["created_at"] = {"$gte": datetime.utcnow() - timedelta(days=7)}
            recent = await documents.count_documents(recent_query)
            
            return {
                "total_documents": total,
                "status_breakdown": status_stats,
                "recent_documents_7d": recent,
            }
        except Exception as e:
            logger.error(f"获取文档统计失败: {e}")
            return {"total_documents": 0, "status_breakdown": {}, "recent_documents_7d": 0}
    
    async def get_stats(self) -> Dict[str, Any]:
        """获取MongoDB统计信息"""
        try:
            # 获取数据库统计
            db_stats = await self.database.command("dbStats")
            
            # 获取集合统计
            collections = await self.database.list_collection_names()
            collection_stats = {}
            
            for collection_name in collections:
                try:
                    collection = self.get_collection(collection_name)
                    count = await collection.count_documents({})
                    collection_stats[collection_name] = count
                except Exception as e:
                    logger.warning(f"获取集合 {collection_name} 统计失败: {e}")
                    collection_stats[collection_name] = 0
            
            return {
                "database_size": db_stats.get("dataSize", 0),
                "storage_size": db_stats.get("storageSize", 0),
                "index_size": db_stats.get("indexSize", 0),
                "collections": collection_stats,
                "total_collections": len(collections)
            }
        except Exception as e:
            logger.error(f"获取MongoDB统计失败: {e}")
            return {
                "database_size": 0,
                "storage_size": 0,
                "index_size": 0,
                "collections": {},
                "total_collections": 0
            }